#自动切图

import glob
import os
import shutil
import xml.etree.cElementTree as et
import cv2
import fire
import numpy as np
import sys
import re
from enum import Enum
import time
import logging
from scipy import signal
from scipy.signal import argrelmin, argrelmax
from scipy.spatial import distance as dist
import matplotlib.pyplot as plt
import xml.dom.minidom as minidom
import pandas as pd



def get_cut_line_hor(img, col_lines, row_lines, pos):
    '''
    寻找切线。
    '''
    thre = 10
    horizontal = np.sum(img[row_lines[0]:row_lines[-1], :], axis=0)
    ht = np.append([0], horizontal[:-1])
    hg = np.subtract(horizontal, ht)
    col_lines_ret = []
    for i, col_line in enumerate(col_lines):
        # print(col_line - thre*5, col_line + thre)
        xend = np.argmax(hg[max(col_line - thre*5, 0): col_line + thre]) + col_line - thre*5
        xstart = np.argmin(hg[max(col_line - thre*5, 0): col_line + thre]) + col_line - thre*5
        if pos[1] in ['1', '3'] and i == 0:
            col_lines_ret.append(xstart)
        else:
            col_lines_ret.append(xend)
    return col_lines_ret


def get_cut_line_ver(img, col_lines, row_lines):
    '''
    寻找切线。
    '''
    hh, ww = img.shape[:2]
    thre = 30
    vertical = np.sum(img[:, col_lines[0]:col_lines[-1]], axis=1)
    # 电池片水平栅格线有可能干扰到电池片的水平切割，特别处理下
    conv_width_half = hh // (2 * 90)
    vertical = np.convolve(vertical, np.ones(
        conv_width_half*2 + 1).astype('int'))[conv_width_half:-1*conv_width_half]
    # 计算梯度
    vt = np.append([0], vertical[:-1])
    vg = np.subtract(vertical, vt)
    row_lines_ret = []
    for row_line in row_lines:
        yend = np.argmax(vg[row_line - thre: row_line + thre]) + row_line - thre
        ystart = np.argmin(vg[row_line - thre: row_line + thre]) + row_line - thre
        row_lines_ret.append(ystart)
    return row_lines_ret


def chuan_cut(img_name, img, bol, col_lines, row_lines, cols):
    chuan_dic = {}
    hh, ww = img.shape[:2]
    col_width = (col_lines[-1] - col_lines[0]) // cols
    for i in range(2):
        if bol[i]:
            top_value = max(row_lines[i]-150, 0)
            down_value = min(row_lines[i]+150, hh)
            for j in range(cols):
                left_value = col_lines[0]+col_width*j
                right_value = min(left_value+col_width, ww)
                save_name = img_name + '_' + str(i) + '_' + str(j+1)
                chuan_dic[save_name] = img[top_value:down_value, left_value:right_value, :]
                # print(chuan_dic[save_name].shape)
    return chuan_dic


def pian_cut(img_name, img, bol, col_lines, row_lines, rows):
    pian_dic = {}
    hh, ww = img.shape[:2]
    row_width = (row_lines[-1] - row_lines[0]) // rows
    for i in range(2):
        if bol[i+2]:
            left_value = max(col_lines[i]-280, 0)
            right_value = min(col_lines[i]+280, ww)
            for j in range(rows):
                top_value = row_lines[0]+row_width*j
                down_value = min(top_value+row_width, hh)
                save_name = img_name + '_' + str(i+1) + '_' + str(j)
                pian_dic[save_name] = np.rot90(img[top_value:down_value, left_value:right_value, :])
                # print(pian_dic[save_name].shape)
    return pian_dic

def slice_get(img_name, img, bol, cols, rows, pad, pos):
    left_lin = int(cols * 375 * pad)
    right_lin = int(cols * 375 * (pad + 1))
    top_lin = int(rows * 750 * pad)
    bottom_lin = int(rows * 750 * (pad + 1))
    t = img[:, :, 0]
    col_lines = get_cut_line_hor(t, (left_lin, right_lin), (top_lin, bottom_lin), pos)
    row_lines = get_cut_line_ver(t, (left_lin, right_lin), (top_lin, bottom_lin))
    dic = {}
    dic['chuan'] = chuan_cut(img_name, img, bol, col_lines, row_lines, cols)
    dic['pian'] = pian_cut(img_name, img, bol, col_lines, row_lines, rows)
    return dic


def check_dir(dir_path):
    '''
    检测路径，如果有路径提示是否删除，无路径新建路径
    '''
    if os.path.exists(dir_path):
        print('!!!dir {} exists'.format(dir_path))
        comfirm = input("!!!remove dir, yes or mix? [y or m]: ")
        if comfirm == 'y':
            print('!!!{} have been removed!!!'.format(dir_path))
            shutil.rmtree(dir_path)
            os.makedirs(dir_path)
        elif comfirm != 'm':
            print('did not remove!')
            sys.exit(0)
    else:
        os.makedirs(dir_path)


def slice_config(pos):
    if pos[0] == 'A':
        top = False
        bottom = True
    elif pos[0] == 'B':
        top = True
        bottom = True
    else:
        top = True
        bottom = False
    if pos[1] in ['1', '3']:
        left = False
        right = True
    else:
        left = True
        right = False
    return top, bottom, left, right


def main(img_dir, 
         save_dir, 
         col_list=[4, 5, 5, 4], 
         row_list=[2, 2, 2], 
         rflip=True, 
         cflip=True, 
         pad=0.2):
    check_dir(save_dir)
    chuan_save_dir = os.path.join(save_dir, 'chuan')
    pian_save_dir = os.path.join(save_dir, 'pian')
    check_dir(chuan_save_dir)
    check_dir(pian_save_dir)
    row_name = ['A', 'B', 'C']
    img_files = glob.glob(os.path.join(img_dir, '*.jpg'))
    il = len(img_files)
    for i, img_file in enumerate(img_files):
        print('<==================={}%====================>'.format(
            (i + 1) * 100 / il))
        img_name = os.path.basename(img_file).split('.')[0]
        print('$$ img {} is processing'.format(img_name))
        pos = img_name.split('_')[-1][0:2]
        if pos[0] not in row_name:
            pos = int(pos)
            img_row = row_name[pos // 4]
            img_col = str(pos % 4 + 1)
            pos = img_row + img_col
        if cflip:
            pos = pos[0] + str(5 - int(pos[1]))
        if rflip:
            pos = row_name[2 - row_name.index(pos[0])] + pos[1]
        cols = col_list[int(pos[1]) - 1]
        rows = 2
        bol = slice_config(pos)
        img = cv2.imdecode(np.fromfile(img_file, dtype=np.uint8), 1)
        dic = slice_get(img_name, img, bol, cols, rows, pad, pos)
        for clas in dic:
            clas_save_dir = os.path.join(save_dir, clas)
            for save_name in dic[clas]:
                save_file = os.path.join(clas_save_dir, save_name + '.jpg')
                cv2.imwrite(save_file, dic[clas][save_name])


if __name__ == '__main__':
    fire.Fire(main)


